from importlib import import_module
from typing import Type, List, Optional, Dict, Any
import yaml

from csagent.core.config.schema import AgentConfig, NodeConfig, EdgeCondition
from csagent.core.context import AgentState
from csagent.core.node.base_node import BaseNode
from csagent.core.node.start_node import StartNode, build_start_node
from csagent.core.node.sub_agent_node import SubAgentNode
from csagent.core.config.loader import load_agent_config
from csagent.core.constants import START_NODE

from langgraph.graph.state import StateGraph, START

def build_subagent_node(node_config: NodeConfig) -> SubAgentNode:
    sub_agent_config = load_agent_config(node_config.conf)
    sub_builder = GraphBuilder(sub_agent_config)
    sub_graph = sub_builder.build()
    node_config = NodeConfig(name=node_config.name)
    sub_agent_node = SubAgentNode(node_config)
    sub_agent_node.graph = sub_graph
    return sub_agent_node

class GraphBuilder:
    def __init__(self, config: AgentConfig):
        self.config = config
        self.graph = StateGraph(AgentState)
        self.nodes = {}  # 存储节点名称到节点的映射
    
    def build(self) -> StateGraph:
        # 1. 添加所有节点
        for node_config in self.config.graph.node:
            self._add_node(node_config)
        # 添加一个虚拟的start节点
        self._add_start_node()
        
        # 2. 添加边
        for source, edges in self.config.graph.edge.items():
            # 将edge info补充到node中
            if source in self.nodes:
                self.nodes[source].edge_condition = edges
            for edge in edges:
                self._add_edge(source, edge)

        return self.graph.compile()
    
    def _add_start_node(self):
        start_node = build_start_node()
        start_node.agent_name = self.config.name
        self.nodes[START_NODE] = start_node
        self.graph.add_node(START_NODE, start_node.run)
        self.graph.add_edge(START, START_NODE)

    def _add_node(self, node_config: NodeConfig):
        """添加单个节点到图中"""
        if node_config.type == 'local_agent':
            node = build_subagent_node(node_config)
        else:       
            # 动态加载自定义节点
            node_class = self._load_node_class(node_config.instance)
        
            # 初始化节点
            node = node_class(node_config)
        node.agent_name = self.config.name
        self.nodes[node_config.name] = node
        self.graph.add_node(node_config.name, node.run)
    
    def _add_edge(self, source: str, edge: EdgeCondition):
        """添加边到图中"""
        pass
        

    def _load_node_class(self, class_path: str) -> Type[BaseNode]:
        """动态加载节点类"""
        module_path, class_name = class_path.rsplit(".", 1)
        module = import_module(module_path)
        return getattr(module, class_name)
    
    def _load_node_config(self, config_path: Optional[str]) -> Dict:
        """加载节点配置"""
        if not config_path:
            return {}
        # 实现配置加载逻辑
        with open(config_path, 'r') as file:
            config = yaml.safe_load(file)
            return config